from gxl_ai_utils.utils import utils_file

original_path = "/home/work_nfs9/zxzhao/workspace/OSUM_emotion_dialog/output_data/jiaofu_jsonl/merged_result_0616.jsonl"
output_path = "/home/work_nfs9/asr_data/data/osum_data/yinshi_emotion_chat/data.jsonl"

reference_path = "/home/work_nfs9/asr_data/data/osum_data/yinshi_emotion_chat/reference_data/reference.jsonl"

data_list = utils_file.load_dict_list_from_jsonl(original_path)
reference_data_list = utils_file.load_dict_list_from_jsonl(reference_path)

reference_dict = {}

for data in reference_data_list:
    reference_dict[data['key']] = data

new_data_list = []

for data in utils_file.tqdm(data_list, desc="Processing data", total=len(data_list)):
    if data['key'] not in reference_dict:
        new_data_list.append(data)
        continue
    reference_dict_item = reference_dict[data['key']]
    q_emotion = reference_dict_item['emotion']
    a_emotion = reference_dict_item['content_emo']
    if data.get('extra', None) is None:
        print(f"Error: extra is not a dict for key {data}")
        continue
    if not isinstance(data['extra'], dict):
        print(f"Error: extra is not a dict for key {data}")
        continue
    data['extra']['q_emotion'] = "<"+q_emotion+">"
    data['extra']['a_emotion'] = "<"+a_emotion+">"
    new_data_list.append(data)

utils_file.write_dict_list_to_jsonl(new_data_list, output_path)



# 先统计reference data
# emotion_dict = {}
# for data in reference_data_list:
#     emotion = data['emotion']
#     if emotion not in emotion_dict:
#         emotion_dict[emotion] = 0
#     emotion_dict[emotion] += 1
# print(emotion_dict)
# {'NEUTRAL': 55738, 'ANGER': 6665, 'FEAR': 6935, 'SAD': 7189, 'DISGUST': 6787, 'SURPRISE': 6921, 'HAPPY': 7199}